import torch
import torch.nn as nn
import torch.optim as optim

from VIT import VisionTransformer

from torchvision import datasets, transforms
# python -m visdom.server
from visdom import Visdom


"""
VIT实验
1. MINST数据集
2. 训练回合数10
3. 训练时长比较短，最终达到90%的准确率，感觉自己做的模型还是有点偏差
"""


# load dataset
def load_dataset():
    batch_size = 100
    # 读取数据集
    train_db = datasets.MNIST('data', train=True, download=True,
                              transform=transforms.Compose([
                                  transforms.ToTensor(),
                                  transforms.Normalize((0.1307,), (0.3081,))
                              ]))
    train_loader = torch.utils.data.DataLoader(
        train_db,
        batch_size=batch_size, shuffle=True)

    test_db = datasets.MNIST('data', train=False, transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ]))
    test_loader = torch.utils.data.DataLoader(
        test_db,
        batch_size=batch_size, shuffle=True)
    # 分割出val数据集
    print('train:', len(train_db), 'test:', len(test_db))
    train_db, val_db = torch.utils.data.random_split(train_db, [50000, 10000])
    print('db1:', len(train_db), 'db2:', len(val_db))

    train_loader = torch.utils.data.DataLoader(
        train_db,
        batch_size=batch_size, shuffle=True)
    val_loader = torch.utils.data.DataLoader(
        val_db,
        batch_size=batch_size, shuffle=True)

    return train_loader, val_loader, test_loader


if __name__ == '__main__':
    epochs = 10
    # 创建网络
    vit = VisionTransformer(
        in_channels=1,
        patch_size=4,
        img_size=28,
        num_heads=4,
        num_layers=6,
        dropout=0.1,
        expand_forward=4,
        n_classes=10
    )

    device = torch.device('cuda')  # 创建cuda设备
    optimizer = optim.Adam(vit.parameters(), lr=5e-4, weight_decay=0.01)  # 创建优化器
    sched = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.5)  # trick: 学习率衰减
    criteon = nn.CrossEntropyLoss().to(device)  # 创建loss函数
    vit.to(device)

    train_loader, val_loader, test_loader = load_dataset()

    for epoch in range(epochs):
        vit.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)

            logits = vit(data)  # 走一次网络[batch,10]
            loss = criteon(logits, target)  # 计算loss
            optimizer.zero_grad()  # 梯度信息清零
            loss.backward()  # 求解梯度
            optimizer.step()  # 利用梯度更新参数

            if batch_idx % 100 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_loader.dataset),
                           100. * batch_idx / len(train_loader), loss.item()))
        # sched.step()

        # 测试网络准确性，一次epoch测试一次
        vit.eval()  # 测试模式
        test_loss = 0
        correct = 0
        for data, target in val_loader:
            data, target = data.to(device), target.cuda()  # torch.Size([28*28]) torch.Size([])
            logits = vit(data)  # 一样的走一遍神经网络 [1,10]
            test_loss += criteon(logits, target).item()  # 累加loss

            pred = logits.argmax(dim=1)  # 得到最大值所在的索引
            correct += pred.eq(target).float().sum().item()  # 累计预测正确的数量

        acc = correct / len(val_loader.dataset)
        print("train accurancy: {}".format(acc))
